package net.kehuilai.tools;

import com.amazonaws.auth.AWSCredentials;
import com.amazonaws.auth.profile.ProfileCredentialsProvider;
import com.amazonaws.regions.Region;
import com.amazonaws.regions.Regions;
import com.amazonaws.services.sns.AmazonSNS;
import com.amazonaws.services.sns.AmazonSNSClient;
import com.amazonaws.services.sns.model.ListTopicsResult;
import com.amazonaws.services.sns.model.PublishRequest;
import com.amazonaws.services.sns.model.Topic;
import com.amazonaws.services.sqs.AmazonSQS;
import com.amazonaws.services.sqs.AmazonSQSClient;
import com.amazonaws.services.sqs.model.SendMessageRequest;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * Created by genghz on 16/11/9.
 */
public class AwsMqSender {

    private static final Logger log = LoggerFactory.getLogger(AwsMqSender.class);
    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    private static final String env = System.getenv().get("ENV");
    private static final AWSCredentials credentials = new ProfileCredentialsProvider().getCredentials();

    private AmazonSNS snsClient;
    private AmazonSQS sqsClient;


    private ConcurrentMap<String, String> topic2arn = new ConcurrentHashMap<String, String>();
    private ConcurrentMap<String, String> queue2url = new ConcurrentHashMap<String, String>();

    public AwsMqSender() {
        this(Regions.CN_NORTH_1);
    }

    public AwsMqSender(Regions regions) {

        Region region = Region.getRegion(regions);

        AmazonSQS sqsClient = new AmazonSQSClient(credentials);
        sqsClient.setRegion(region);

        AmazonSNS snsClient = new AmazonSNSClient(credentials);
        snsClient.setRegion(region);

        this.snsClient = snsClient;
        this.sqsClient = sqsClient;
    }

    public void push(String name, Object o) {
        pushTopic(name, o, 0);
    }

    public void push(String name, Object o, int delaySeconds) {
        pushTopic(name, o, delaySeconds);
    }

    private void pushTopic(String topicName, Object o, int delaySeconds) {

        String topicArn = getTopicArn(topicName);
        PublishRequest request = new PublishRequest(topicArn, pojo2json(o));
        request.addMessageAttributesEntry("_type_", setSnsMessageAttributeValue("String", o.getClass().getName()));
        if (delaySeconds > 0)
            request.addMessageAttributesEntry("_delay_", setSnsMessageAttributeValue("String", String.valueOf(delaySeconds)));

        snsClient.publish(request);
    }

    private void pushQueue(String queueName, Object o, int delaySeconds) {

        String queueUrl = getQueueUrl(queueName);

        SendMessageRequest request = new SendMessageRequest(queueUrl, pojo2json(o));
        request.addMessageAttributesEntry("_type_", setSqsMessageAttributeValue("String", o.getClass().getName()));
        if (delaySeconds > 0)
            request.setDelaySeconds(delaySeconds);

        sqsClient.sendMessage(request);
    }

    // tool
    private String getTopicArn(String topicName) {
        String name = getName(topicName);

        String topicArn = this.topic2arn.get(name);
        if (topicArn == null || "".equals(topicArn)) {
            ListTopicsResult result = this.snsClient.listTopics();
            for (Topic topic : result.getTopics()) {
                String arn = topic.getTopicArn();
                String[] arns = arn.split(":");
                this.topic2arn.put(arns[arns.length - 1], arn);
            }
            topicArn = this.topic2arn.get(name);
        }
        return topicArn;
    }

    private String getQueueUrl(String queueName) {
        String name = getName(queueName);

        if (!queue2url.containsKey(name)) {
            queue2url.put(name, this.sqsClient.getQueueUrl(name).getQueueUrl());
        }
        return this.queue2url.get(name);
    }

    private String getName(String name) {
        if ("prod".equals(env)) {
            return name;
        } else {
            return name + "-test";
        }
    }

    private com.amazonaws.services.sqs.model.MessageAttributeValue setSqsMessageAttributeValue(String dataType, String value) {
        return new com.amazonaws.services.sqs.model.MessageAttributeValue()
                .withDataType(dataType)
                .withStringValue(value);
    }

    private com.amazonaws.services.sns.model.MessageAttributeValue setSnsMessageAttributeValue(String dataType, String value) {
        return new com.amazonaws.services.sns.model.MessageAttributeValue()
                .withDataType(dataType)
                .withStringValue(value);
    }

    private String pojo2json(Object o) {
        if (o == null) return null;
        String json = null;
        try {
            json = OBJECT_MAPPER.writeValueAsString(o);
        } catch (JsonProcessingException e) {
            log.warn("pojo2json failed: {}", e.getMessage());
        }
        return json;
    }
}
